from fastapi import FastAPI, UploadFile, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.encoders import jsonable_encoder
from paddleocr import PaddleOCR
from PIL import Image
import io
import numpy as np
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import asyncio
import concurrent.futures
from functools import lru_cache
import logging
from logging.handlers import RotatingFileHandler
import os
import time
import uvicorn
from pydantic import BaseModel
from typing import List, Any
import cv2
import multiprocessing
import psutil
import pillow_avif
import traceback

# 异常和错误响应模型定义
class AppException(Exception):
    def __init__(self, error_code: str, message: str, status_code: int = 400):
        self.error_code = error_code
        self.message = message
        self.status_code = status_code

class ErrorResponse(BaseModel):
    error_code: str
    message: str
    details: Any = None

app = FastAPI()

# 配置日志
def setup_logger():
    log_dir = 'logs'
    os.makedirs(log_dir, exist_ok=True)
    log_file = os.path.join(log_dir, 'app.log')

    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    file_handler = RotatingFileHandler(log_file, maxBytes=10000000, backupCount=5)
    file_handler.setFormatter(formatter)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)

    logger = logging.getLogger("app")
    logger.setLevel(logging.INFO)
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger

logger = setup_logger()

# 初始化PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang="ch")

# 设置线程池大小为CPU核心数的2倍
thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=multiprocessing.cpu_count())

def get_image_with_retry(url, max_retries=3, backoff_factor=0.3):
    session = requests.Session()
    retry = Retry(total=max_retries,
                  backoff_factor=backoff_factor,
                  status_forcelist=[420, 429, 500, 502, 503, 504])
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Accept': 'image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8',
        'Accept-Language': 'en-US,en;q=0.9',
        'Referer': 'https://www.google.com/'
    }
    
    try:
        response = session.get(url, headers=headers, timeout=10)
        response.raise_for_status()
        logger.info(f"成功获取图像，URL: {url}, 状态码: {response.status_code}, 内容类型: {response.headers.get('content-type')}")
        return response
    except requests.exceptions.RequestException as e:
        logger.error(f"获取图像失败，URL: {url}, 错误: {str(e)}")
        raise

def log_exception(e: Exception, context: str = ""):
    logger.error(f"{context} - 异常类型: {type(e).__name__}, 异常信息: {str(e)}")
    logger.error(f"异常堆栈: {''.join(traceback.format_tb(e.__traceback__))}")

async def get_image(source):
    logger.info(f"获取图像: {source}")
    start_time = time.time()
    try:
        if isinstance(source, str) and source.startswith(('http://', 'https://')):
            response = get_image_with_retry(source)
            content_type = response.headers.get('content-type', '')
            if 'image' not in content_type:
                raise ValueError(f"无效的内容类型: {content_type}")
            image_data = response.content
        elif isinstance(source, (io.BytesIO, UploadFile)):
            if isinstance(source, UploadFile):
                image_data = await source.read()
            else:
                image_data = source.getvalue()
        else:
            raise ValueError("无效的图片来源")

        img_pil = Image.open(io.BytesIO(image_data))
        img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

        logger.info(f"图像获取成功，格式: {img_pil.format}, 大小: {img_pil.size}, 耗时: {time.time() - start_time:.2f}秒")
        return img
    except Exception as e:
        log_exception(e, "获取图像失败")
        raise AppException(error_code="IMAGE_FETCH_ERROR", message=f"获取图像失败: {str(e)}")

@lru_cache(maxsize=1000)
def cached_ocr(image_hash):
    logger.info(f"执行OCR: {image_hash[:20]}...（哈希前20字节）")
    start_time = time.time()
    
    try:
        img = Image.open(io.BytesIO(image_hash))
        result = ocr.ocr(np.array(img))
        ocr_time = time.time() - start_time
        logger.info(f"OCR完成，识别到{len(result[0])}个文本块，耗时: {ocr_time:.2f}秒")
        return {"success": True, "result": result[0]}
    except Exception as e:
        log_exception(e, "OCR执行失败")
        return {"success": False, "error": str(e)}

async def process_image(image_source):
    logger.info(f"处理图像: {type(image_source)}")
    start_time = time.time()
    try:
        img = await get_image(image_source)
        
        max_size = 1024
        height, width = img.shape[:2]
        if height > max_size or width > max_size:
            scale = max_size / max(height, width)
            img = cv2.resize(img, None, fx=scale, fy=scale)
        
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_pil = Image.fromarray(img_rgb)
        
        img_byte_arr = io.BytesIO()
        img_pil.save(img_byte_arr, format='JPEG', quality=85)
        image_hash = img_byte_arr.getvalue()
        
        cache_info = cached_ocr.cache_info()
        logger.info(f"缓存状态 - 命中: {cache_info.hits}, 未命中: {cache_info.misses}, 当前大小: {cache_info.currsize}, 最大大小: {cache_info.maxsize}")
        
        loop = asyncio.get_running_loop()
        ocr_result = await loop.run_in_executor(thread_pool, cached_ocr, image_hash)
        
        if cached_ocr.cache_info().hits > cache_info.hits:
            logger.info("OCR结果从缓存中获取")
        else:
            logger.info("OCR结果新生成并缓存")
        
        process_time = time.time() - start_time
        logger.info(f"单个图像处理完成，耗时: {process_time:.2f}秒")
        return {"success": True, "result": ocr_result}
    except AppException as e:
        logger.error(f"应用异常: {e.error_code} - {e.message}")
        raise
    except Exception as e:
        log_exception(e, "图像处理失败")
        raise AppException(error_code="IMAGE_PROCESS_ERROR", message=f"图像处理失败: {str(e)}")

@app.on_event("startup")
async def startup_event():
    cache_info = cached_ocr.cache_info()
    logger.info(f"应用启动 - 初始缓存状态: {cache_info}")
    asyncio.create_task(log_cache_status())

async def log_cache_status():
    while True:
        cache_info = cached_ocr.cache_info()
        logger.info(f"定期缓存状态报告: {cache_info}")
        await asyncio.sleep(3600)

class OCRRequest(BaseModel):
    image_url: str

@app.post("/ocr")
async def perform_ocr(request: Request):
    logger.info("接收到OCR请求")
    logger.info(f"请求头: {request.headers}")
    start_time = time.time()
    try:
        content_type = request.headers.get('Content-Type', '')
        logger.info(f"Content-Type: {content_type}")
        
        if content_type.startswith('application/json'):
            body = await request.json()
            logger.info(f"JSON body: {body}")
            image_url = body.get('image_url')
            if not image_url:
                raise HTTPException(status_code=400, detail="JSON请求必须提供image_url")
            result = await process_image(image_url)
        elif content_type.startswith('multipart/form-data'):
            form = await request.form()
            logger.info(f"Form data keys: {form.keys()}")
            if 'image' in form:
                image = await form['image'].read()
                logger.info(f"Received image file: {form['image'].filename}, size: {len(image)} bytes")
                result = await process_image(io.BytesIO(image))
            elif 'image_url' in form:
                image_url = form['image_url']
                result = await process_image(image_url)
            else:
                raise HTTPException(status_code=400, detail="表单数据必须提供image或image_url")
        else:
            raise HTTPException(status_code=400, detail=f"不支持的Content-Type: {content_type}")
        
        logger.info(f"OCR请求处理完成，耗时: {time.time() - start_time:.2f}秒")
        return JSONResponse(content=jsonable_encoder({"success": True, "texts": result}), media_type="application/json; charset=utf-8")
    except AppException as e:
        logger.error(f"OCR请求处理失败 : {e.error_code} - {e.message}")
        raise
    except Exception as e:
        log_exception(e, "OCR请求处理失败")
        raise AppException(error_code="OCR_ERROR", message=f"OCR请求处理失败: {str(e)}")

class BatchOCRRequest(BaseModel):
    image_sources: List[str]

@app.post("/ocr_batch")
async def perform_ocr_batch(request: BatchOCRRequest):
    logger.info("接收到批量OCR请求")
    start_time = time.time()
    try:
        cpu_percent = psutil.cpu_percent()
        memory_percent = psutil.virtual_memory().percent
        logger.info(f"批处理开始 - CPU使用率: {cpu_percent}%, 内存使用率: {memory_percent}%")

        tasks = [process_image(source) for source in request.image_sources]
        results = await asyncio.gather(*tasks)

        cpu_percent = psutil.cpu_percent()
        memory_percent = psutil.virtual_memory().percent
        logger.info(f"批处理结束 - CPU使用率: {cpu_percent}%, 内存使用率: {memory_percent}%")

        total_time = time.time() - start_time
        logger.info(f"批量OCR请求处理完成，共处理{len(results)}个图片，总耗时: {total_time:.2f}秒，平均每张耗时: {total_time/len(results):.2f}秒")
        
        return JSONResponse(
            content=jsonable_encoder({"success": True, "results": results}),
            media_type="application/json; charset=utf-8"
        )
    except AppException as e:
        logger.error(f"批量OCR请求处理失败: {e.error_code} - {e.message}")
        raise
    except Exception as e:
        log_exception(e, "批量OCR请求处理失败")
        raise AppException(error_code="BATCH_OCR_ERROR", message=f"批量OCR请求处理失败: {str(e)}")

@app.exception_handler(AppException)
async def app_exception_handler(request: Request, exc: AppException):
    logger.error(f"应用异常: {exc.error_code} - {exc.message}")
    return JSONResponse(
        status_code=exc.status_code,
        content=jsonable_encoder(ErrorResponse(error_code=exc.error_code, message=exc.message))
    )

@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    log_exception(exc, "未处理的异常")
    return JSONResponse(
        status_code=500,
        content=jsonable_encoder(ErrorResponse(error_code="INTERNAL_ERROR", message="发生了意外错误，请稍后再试。"))
    )

if __name__ == '__main__':
    logger.info("应用启动")
    uvicorn.run(app, host="0.0.0.0", port=5000, workers=4)